"""
python3 script/gen_serendipity_basis.py > sfepy/discrete/fem/_serendipity.py
"""
import sympy as sm

x, y, z = sm.symbols('x y z')

all_bfs = {
    2: {
        1 : [
            0.25 * (1 - x) * (1 - y),
            0.25 * (1 + x) * (1 - y),
            0.25 * (1 + x) * (1 + y),
            0.25 * (1 - x) * (1 + y),
        ],
        2 : [
            -0.25 * (1 - x) * (1 - y) * (1 + x + y),
            -0.25 * (1 + x) * (1 - y) * (1 - x + y),
            -0.25 * (1 + x) * (1 + y) * (1 - x - y),
            -0.25 * (1 - x) * (1 + y) * (1 + x - y),
            0.5 * (1 - x**2) * (1 - y),
            0.5 * (1 + x) * (1 - y**2),
            0.5 * (1 - x**2) * (1 + y),
            0.5 * (1 - x) * (1 - y**2),
        ],
        3 : [
            0.03125 * (x - 1) * (y - 1) * (9 * (x**2 + y**2) - 10),
            -0.03125 * (x + 1) * (y - 1) * (9 * (x**2 + y**2) - 10),
            0.03125 * (x + 1) * (y + 1) * (9 * (x**2 + y**2) - 10),
            -0.03125 * (x - 1) * (y + 1) * (9 * (x**2 + y**2) - 10),
            0.28125 * (y - 1) * (-3 * x**3 + x**2 + 3 * x - 1),
            -0.28125 * (y - 1) * (-3 * x**3 - x**2 + 3 * x + 1),
            -0.28125 * (x + 1) * (-3 * y**3 + y**2 + 3 * y - 1),
            0.28125 * (x + 1) * (-3 * y**3 - y**2 + 3 * y + 1),
            0.28125 * (y + 1) * (-3 * x**3 - x**2 + 3 * x + 1),
            -0.28125 * (y + 1) * (-3 * x**3 + x**2 + 3 * x - 1),
            -0.28125 * (x - 1) * (-3 * y**3 - y**2 + 3 * y + 1),
            0.28125 * (x - 1) * (-3 * y**3 + y**2 + 3 * y - 1),
        ],
    },
    3 : {
        1 : [
            0.125 * (1 - x) * (1 - y) * (1 - z),
            0.125 * (1 + x) * (1 - y) * (1 - z),
            0.125 * (1 + x) * (1 + y) * (1 - z),
            0.125 * (1 - x) * (1 + y) * (1 - z),
            0.125 * (1 - x) * (1 - y) * (1 + z),
            0.125 * (1 + x) * (1 - y) * (1 + z),
            0.125 * (1 + x) * (1 + y) * (1 + z),
            0.125 * (1 - x) * (1 + y) * (1 + z),
        ],
        2 : [
            -0.125 * (1 - x) * (1 - y) * (1 - z) * (x + y + z + 2),
            -0.125 * (1 + x) * (1 - y) * (1 - z) * (-x + y + z + 2),
            -0.125 * (1 + x) * (1 + y) * (1 - z) * (-x - y + z + 2),
            -0.125 * (1 - x) * (1 + y) * (1 - z) * (x - y + z + 2),
            -0.125 * (1 - x) * (1 - y) * (1 + z) * (x + y - z + 2),
            -0.125 * (1 + x) * (1 - y) * (1 + z) * (-x + y - z + 2),
            -0.125 * (1 + x) * (1 + y) * (1 + z) * (-x - y - z + 2),
            -0.125 * (1 - x) * (1 + y) * (1 + z) * (x - y - z + 2),
            0.25 * (1 - x) * (1 + x) * (1 - y) * (1 - z),
            0.25 * (1 - y) * (1 + y) * (1 + x) * (1 - z),
            0.25 * (1 - x) * (1 + x) * (1 + y) * (1 - z),
            0.25 * (1 - y) * (1 + y) * (1 - x) * (1 - z),
            0.25 * (1 - x) * (1 + x) * (1 - y) * (1 + z),
            0.25 * (1 - y) * (1 + y) * (1 + x) * (1 + z),
            0.25 * (1 - x) * (1 + x) * (1 + y) * (1 + z),
            0.25 * (1 - y) * (1 + y) * (1 - x) * (1 + z),
            0.25 * (1 - z) * (1 + z) * (1 - x) * (1 - y),
            0.25 * (1 - z) * (1 + z) * (1 + x) * (1 - y),
            0.25 * (1 - z) * (1 + z) * (1 + x) * (1 + y),
            0.25 * (1 - z) * (1 + z) * (1 - x) * (1 + y),
        ],
        3 : [
            0.015625 * (1 - x) * (1 - y) * (1 - z) * (9 * (x**2 + y**2 + z**2) - 19.),
            0.015625 * (1 + x) * (1 - y) * (1 - z) * (9 * (x**2 + y**2 + z**2) - 19.),
            0.015625 * (1 + x) * (1 + y) * (1 - z) * (9 * (x**2 + y**2 + z**2) - 19.),
            0.015625 * (1 - x) * (1 + y) * (1 - z) * (9 * (x**2 + y**2 + z**2) - 19.),
            0.015625 * (1 - x) * (1 - y) * (1 + z) * (9 * (x**2 + y**2 + z**2) - 19.),
            0.015625 * (1 + x) * (1 - y) * (1 + z) * (9 * (x**2 + y**2 + z**2) - 19.),
            0.015625 * (1 + x) * (1 + y) * (1 + z) * (9 * (x**2 + y**2 + z**2) - 19.),
            0.015625 * (1 - x) * (1 + y) * (1 + z) * (9 * (x**2 + y**2 + z**2) - 19.),

            -0.140625 * (1 - y) * (1 - z) * (-3 * x**3 + x**2 + 3 * x - 1),
            0.140625 * (1 - y) * (1 - z) * (-3 * x**3 - x**2 + 3 * x + 1),
            -0.140625 * (1 + x) * (1 - z) * (-3 * y**3 + y**2 + 3 * y - 1),
            0.140625 * (1 + x) * (1 - z) * (-3 * y**3 - y**2 + 3 * y + 1),

            0.140625 * (1 + y) * (1 - z) * (-3 * x**3 - x**2 + 3 * x + 1),
            -0.140625 * (1 + y) * (1 - z) * (-3 * x**3 + x**2 + 3 * x - 1),
            0.140625 * (1 - x) * (1 - z) * (-3 * y**3 - y**2 + 3 * y + 1),
            -0.140625 * (1 - x) * (1 - z) * (-3 * y**3 + y**2 + 3 * y - 1),

            -0.140625 * (1 - y) * (1 + z) * (-3 * x**3 + x**2 + 3 * x - 1),
            0.140625 * (1 - y) * (1 + z) * (-3 * x**3 - x**2 + 3 * x + 1),
            -0.140625 * (1 + x) * (1 + z) * (-3 * y**3 + y**2 + 3 * y - 1),
            0.140625 * (1 + x) * (1 + z) * (-3 * y**3 - y**2 + 3 * y + 1),

            0.140625 * (1 + y) * (1 + z) * (-3 * x**3 - x**2 + 3 * x + 1),
            -0.140625 * (1 + y) * (1 + z) * (-3 * x**3 + x**2 + 3 * x - 1),
            0.140625 * (1 - x) * (1 + z) * (-3 * y**3 - y**2 + 3 * y + 1),
            -0.140625 * (1 - x) * (1 + z) * (-3 * y**3 + y**2 + 3 * y - 1),

            -0.140625 * (1 - x) * (1 - y) * (-3 * z**3 + z**2 + 3 * z - 1),
            0.140625 * (1 - x) * (1 - y) * (-3 * z**3 - z**2 + 3 * z + 1),
            -0.140625 * (1 + x) * (1 - y) * (-3 * z**3 + z**2 + 3 * z - 1),
            0.140625 * (1 + x) * (1 - y) * (-3 * z**3 - z**2 + 3 * z + 1),

            -0.140625 * (1 + x) * (1 + y) * (-3 * z**3 + z**2 + 3 * z - 1),
            0.140625 * (1 + x) * (1 + y) * (-3 * z**3 - z**2 + 3 * z + 1),
            -0.140625 * (1 - x) * (1 + y) * (-3 * z**3 + z**2 + 3 * z - 1),
            0.140625 * (1 - x) * (1 + y) * (-3 * z**3 - z**2 + 3 * z + 1),
        ],
    }
}

header = """
import sympy as sm

x, y, z = sm.symbols('x y z')
"""

def main():
    print(header.lstrip())

    print('all_bfs = {')
    for dim, _dbfs in all_bfs.items():
        print('    {} : {{'.format(dim))
        for order, _bfs in _dbfs.items():
            print('        {} : ['.format(order))
            vs = [x, y, z][:dim]
            bfs = [
                sm.horner(sm.simplify(
                    bf.subs({x : -1 + 2 * x, y : -1 + 2 * y, z : -1 + 2 * z}))
                ) for bf in _bfs
            ]
            bfgs = [[sm.horner(sm.simplify(bf.diff(v)))
                      for v in vs] for bf in bfs]

            print('                [')
            for bf in bfs:
                print((' ' * 16), str(bf).replace('1.0*', ''), ',')
            print('                ],')
            print('                [')
            for bfg in bfgs:
                print((' ' * 16), str(bfg).replace('1.0*', ''), ',')
            print('                ],')

            print('        ],')
        print('    },')
    print('}')

if __name__ == '__main__':
    main()
